# Make an image of a single dataset.

import os
import sys
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import readData
import random

xLims = (-2.0,2.0)
yLims = (-2.0,2.0)

# Given a file containing a data set, make some plots.
def plotData(dataFn):
   data = readData.readData(dataFn)

   xPts = []
   yPts = []

   for r in data:
      xPts.append(r[1])
      yPts.append(r[2])

   fig = plt.figure()
   ax = fig.add_subplot(111)

   ax.plot(xPts, yPts, 'ro', markersize = 2)
   ax.set_xlim(xLims)
   ax.set_ylim(yLims)

   saveFn = dataFn[dataFn.rfind('DD'):]
   print saveFn
   plt.savefig("../plots/dataPlot_" + saveFn + ".png")
   plt.close()

# Plot the positions of the bodies as a set of orbits.
def plotLines():

   xPts = []
   yPts = []

   data = readData.readAllData("../parameters/MWBH.txt")
   fig = plt.figure()
   ax = fig.add_subplot(111)

   for r in data:
      xPts = []
      yPts = []

      for p in r:
         xPts.append(p[1])
         yPts.append(p[2])

      col = [random.random(), random.random(), random.random()]
      ax.plot(xPts, yPts, 'o', markerfacecolor = col, markersize = 3)


  # ax.plot(xPts, yPts, 'ko', markersize = 3)
#   ax.set_xlim(xLims)
#   ax.set_ylim(yLims)

   plt.savefig("../plots/lines.png")

plotLines()

# Plots potential, kinetic, and total energy as a function of time.
def plotEnergy(paramFn):
   KE = []
   PE = []
   TE = []
   t = []

   data = readData.readAllData(paramFn)
   outNum = 0

   G = 1.0
   G_cgs = 1.0

   MassUnits = 1.0
   TimeUnits = 1.0
   LengthUnits = 1.0

   inf = open(paramFn)

   for line in inf:
      line = (line.split("//"))[0]
      line = (line.split("#"))[0]
      ele = line.split("=")

      if ele[0].strip() == "G_cgs":
         G_cgs = float(ele[1].strip())

      if ele[0].strip() == "MassUnits":
         MassUnits = float(ele[1].strip())

      if ele[0].strip() == "TimeUnits":
         TimeUnits = float(ele[1].strip())

      if ele[0].strip() == "LengthUnits":
         LengthUnits= float(ele[1].strip())

   print MassUnits, TimeUnits, LengthUnits, G_cgs
   G = G_cgs / (LengthUnits**3.0 / (MassUnits * TimeUnits**2.0))

   for DD in data:
     # Compute the kinetic energy.
     k = 0.0

     for r in DD:
        k = k + r[0] * (pow(r[4],2.0) + pow(r[5],2.0) + pow(r[6],2.0)) * 0.5

     KE.append(k)

     # Compute the potential energy
     u = 0.0

     for i in range(len(DD)):
        for j in range(i+1,len(DD)):
	   dist = np.sqrt(pow(DD[i][1] - DD[j][1], 2.0)
	         + pow(DD[i][2] - DD[j][2], 2.0)
	         + pow(DD[i][3] - DD[j][3], 2.0));

	   u = u + G * DD[i][0] * DD[j][0] / dist;

     PE.append(u);
     TE.append(k-u);
     print k-u

     t.append(outNum);
     outNum = outNum + 1

   fig = plt.figure()
   ax = fig.add_subplot(111)

   ax.plot(t, KE, 'blue', linewidth = 1, label = "Kinetic")
   ax.plot(t, PE, 'green', linewidth = 1, label = "Potential")
   ax.plot(t, TE, 'black', linewidth= 2, label = "Total")

   ax.legend()
   plt.savefig("energy.png")
